/*================================================================
*   Copyright (C) 2024 baichao All rights reserved.
*
*   文件名称：mq_notify_select.c
*   创 建 者：baichao
*   创建日期：2024年06月22日
*   描    述：
*
================================================================*/

#include <fcntl.h>
#include <mqueue.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <sys/wait.h>
#include <unistd.h>
#include <errno.h>

#define QUEUE_NAME "/test_queue"
#define MAX_SIZE 1024
#define MSG_STOP "exit"

mqd_t mq;

// 消息队列属性
struct mq_attr attr = {
    .mq_flags = 0, .mq_maxmsg = 10, .mq_msgsize = MAX_SIZE, .mq_curmsgs = 0};

// 清理消息队列
void cleanup() { mq_unlink(QUEUE_NAME); }

// 生产者函数
void producer() {
  mqd_t mq;
  char buffer[MAX_SIZE];
  int count = 0;
  unsigned int priority;

  // 打开消息队列
  mq = mq_open(QUEUE_NAME, O_WRONLY);
  if (mq == (mqd_t)-1) {
    perror("Producer: mq_open");
    exit(1);
  }

  while (1) {
    priority = rand() % 5;  // 随机优先级 0-4
    snprintf(buffer, MAX_SIZE, "Message %d with priority %d", count++,
             priority);
    printf("Producer: Sending message: %s\n", buffer);
    if (mq_send(mq, buffer, strlen(buffer) + 1, priority) == -1) {
      perror("Producer: mq_send");
      exit(1);
    }
    sleep(1);  // 模拟生产者的工作延迟

    // 发送停止信号
    if (count > 5) {
      mq_send(mq, MSG_STOP, strlen(MSG_STOP) + 1, 0);
      break;
    }
  }

  // 关闭消息队列
  mq_close(mq);
}

void handle_notification(int sig);
// 消费者函数
void consumer() {
  fd_set rfds;
  char buffer[MAX_SIZE + 1];
  unsigned int priority;
  struct mq_attr attr;
  int retval;

  // 打开消息队列
  mq = mq_open(QUEUE_NAME, O_RDONLY | O_NONBLOCK);
  if (mq == (mqd_t)-1) {
    perror("Consumer: mq_open");
    exit(1);
  }

  // 注册消息队列通知
  struct sigevent sev;
  sev.sigev_notify = SIGEV_SIGNAL;
  sev.sigev_signo = SIGUSR1;

  if (mq_notify(mq, &sev) == -1) {
    perror("Consumer: mq_notify");
    exit(1);
  }

  // 设置 SIGUSR1 信号的处理函数
  signal(SIGUSR1, handle_notification);

  // 无限循环等待信号
  while (1) {
    pause();  // 等待信号到达
  }

  // 关闭消息队列
  mq_close(mq);
}

// 信号处理程序，用于通知消息到达
void handle_notification(int sig) {
  char buffer[MAX_SIZE + 1];
  unsigned int priority;
  struct mq_attr attr;
  fd_set rfds;
  int retval;

  // 获取消息队列属性
  if (mq_getattr(mq, &attr) == -1) {
    perror("Consumer: mq_getattr");
    exit(1);
  }

  // 重新注册通知
  struct sigevent sev;
  sev.sigev_notify = SIGEV_SIGNAL;
  sev.sigev_signo = SIGUSR1;

  if (mq_notify(mq, &sev) == -1) {
    perror("Consumer: mq_notify");
    exit(1);
  }

  while (1) {
    // 初始化 fd_set
    FD_ZERO(&rfds);
    FD_SET(mq, &rfds);

    // 使用 select 等待消息
    retval = select(mq + 1, &rfds, NULL, NULL, NULL);
    if (retval == -1) {
      perror("Consumer: select");
      exit(1);
    }

    if (FD_ISSET(mq, &rfds)) {
      // 接收所有可用消息
      while (mq_receive(mq, buffer, MAX_SIZE, &priority) != -1) {
        buffer[MAX_SIZE] = '\0';
        printf("Consumer: Received message: %s with priority %u\n", buffer,
               priority);

        // 检查是否收到停止信号
        if (strcmp(buffer, MSG_STOP) == 0) {
          printf("Consumer: Received stop signal.\n");
          exit(0);
        }
      }

      if (errno != EAGAIN) {
        perror("Consumer: mq_receive");
        exit(1);
      }
    }
  }
}

int main() {
  // 注册退出时的清理函数
  atexit(cleanup);

  // 创建消息队列
  mq = mq_open(QUEUE_NAME, O_CREAT | O_RDWR, 0644, &attr);
  if (mq == (mqd_t)-1) {
    perror("Main: mq_open");
    exit(1);
  }

  // 创建子进程
  pid_t pid = fork();

  if (pid == -1) {
    perror("Main: fork");
    exit(1);
  } else if (pid == 0) {
    // 子进程作为生产者
    producer();
  } else {
    // 父进程作为消费者
    consumer();

    // 等待子进程结束
    wait(NULL);
  }

  return 0;
}
